{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *\n",
    "from fastai.callbacks.hooks import *\n",
    "from fastai.utils.mem import *\n",
    "from helper import *\n",
    "from torchsummary import summary\n",
    "import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['resnet10', 'resnet14', 'resnet18', 'resnet20', 'resnet26', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'resnext101_32x16d', 'resnext101_32x32d', 'resnext101_32x48d', 'dpn68', 'dpn68b', 'dpn92', 'dpn98', 'dpn107', 'dpn131', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn', 'senet154', 'se_resnet50', 'se_resnet101', 'se_resnet152', 'se_resnext50_32x4d', 'se_resnext101_32x4d', 'densenet121', 'densenet169', 'densenet201', 'densenet161', 'inceptionresnetv2', 'inceptionv4', 'efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2', 'efficientnet-b3', 'efficientnet-b4', 'efficientnet-b5', 'efficientnet-b6', 'efficientnet-b7', 'mobilenet_v2', 'xception']\n"
     ]
    }
   ],
   "source": [
    "encoders = models.encoders.get_encoder_names()\n",
    "print(encoders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 240, 320]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 240, 320]             128\n",
      "              ReLU-3         [-1, 64, 240, 320]               0\n",
      "         MaxPool2d-4         [-1, 64, 120, 160]               0\n",
      "            Conv2d-5         [-1, 64, 120, 160]          36,864\n",
      "       BatchNorm2d-6         [-1, 64, 120, 160]             128\n",
      "              ReLU-7         [-1, 64, 120, 160]               0\n",
      "            Conv2d-8         [-1, 64, 120, 160]          36,864\n",
      "       BatchNorm2d-9         [-1, 64, 120, 160]             128\n",
      "             ReLU-10         [-1, 64, 120, 160]               0\n",
      "       BasicBlock-11         [-1, 64, 120, 160]               0\n",
      "           Conv2d-12          [-1, 128, 60, 80]          73,728\n",
      "      BatchNorm2d-13          [-1, 128, 60, 80]             256\n",
      "             ReLU-14          [-1, 128, 60, 80]               0\n",
      "           Conv2d-15          [-1, 128, 60, 80]         147,456\n",
      "      BatchNorm2d-16          [-1, 128, 60, 80]             256\n",
      "           Conv2d-17          [-1, 128, 60, 80]           8,192\n",
      "      BatchNorm2d-18          [-1, 128, 60, 80]             256\n",
      "             ReLU-19          [-1, 128, 60, 80]               0\n",
      "       BasicBlock-20          [-1, 128, 60, 80]               0\n",
      "           Conv2d-21          [-1, 256, 30, 40]         294,912\n",
      "      BatchNorm2d-22          [-1, 256, 30, 40]             512\n",
      "             ReLU-23          [-1, 256, 30, 40]               0\n",
      "           Conv2d-24          [-1, 256, 30, 40]         589,824\n",
      "      BatchNorm2d-25          [-1, 256, 30, 40]             512\n",
      "           Conv2d-26          [-1, 256, 30, 40]          32,768\n",
      "      BatchNorm2d-27          [-1, 256, 30, 40]             512\n",
      "             ReLU-28          [-1, 256, 30, 40]               0\n",
      "       BasicBlock-29          [-1, 256, 30, 40]               0\n",
      "           Conv2d-30          [-1, 512, 15, 20]       1,179,648\n",
      "      BatchNorm2d-31          [-1, 512, 15, 20]           1,024\n",
      "             ReLU-32          [-1, 512, 15, 20]               0\n",
      "           Conv2d-33          [-1, 512, 15, 20]       2,359,296\n",
      "      BatchNorm2d-34          [-1, 512, 15, 20]           1,024\n",
      "           Conv2d-35          [-1, 512, 15, 20]         131,072\n",
      "      BatchNorm2d-36          [-1, 512, 15, 20]           1,024\n",
      "             ReLU-37          [-1, 512, 15, 20]               0\n",
      "       BasicBlock-38          [-1, 512, 15, 20]               0\n",
      "    ResNetEncoder-39  [[-1, 3, 480, 640], [-1, 64, 240, 320], [-1, 64, 120, 160], [-1, 128, 60, 80], [-1, 256, 30, 40], [-1, 512, 15, 20]]               0\n",
      "         Identity-40          [-1, 512, 15, 20]               0\n",
      "         Identity-41          [-1, 768, 30, 40]               0\n",
      "        Attention-42          [-1, 768, 30, 40]               0\n",
      "           Conv2d-43          [-1, 256, 30, 40]       1,769,472\n",
      "      BatchNorm2d-44          [-1, 256, 30, 40]             512\n",
      "             ReLU-45          [-1, 256, 30, 40]               0\n",
      "           Conv2d-46          [-1, 256, 30, 40]         589,824\n",
      "      BatchNorm2d-47          [-1, 256, 30, 40]             512\n",
      "             ReLU-48          [-1, 256, 30, 40]               0\n",
      "         Identity-49          [-1, 256, 30, 40]               0\n",
      "        Attention-50          [-1, 256, 30, 40]               0\n",
      "     DecoderBlock-51          [-1, 256, 30, 40]               0\n",
      "         Identity-52          [-1, 384, 60, 80]               0\n",
      "        Attention-53          [-1, 384, 60, 80]               0\n",
      "           Conv2d-54          [-1, 128, 60, 80]         442,368\n",
      "      BatchNorm2d-55          [-1, 128, 60, 80]             256\n",
      "             ReLU-56          [-1, 128, 60, 80]               0\n",
      "           Conv2d-57          [-1, 128, 60, 80]         147,456\n",
      "      BatchNorm2d-58          [-1, 128, 60, 80]             256\n",
      "             ReLU-59          [-1, 128, 60, 80]               0\n",
      "         Identity-60          [-1, 128, 60, 80]               0\n",
      "        Attention-61          [-1, 128, 60, 80]               0\n",
      "     DecoderBlock-62          [-1, 128, 60, 80]               0\n",
      "         Identity-63        [-1, 192, 120, 160]               0\n",
      "        Attention-64        [-1, 192, 120, 160]               0\n",
      "           Conv2d-65         [-1, 64, 120, 160]         110,592\n",
      "      BatchNorm2d-66         [-1, 64, 120, 160]             128\n",
      "             ReLU-67         [-1, 64, 120, 160]               0\n",
      "           Conv2d-68         [-1, 64, 120, 160]          36,864\n",
      "      BatchNorm2d-69         [-1, 64, 120, 160]             128\n",
      "             ReLU-70         [-1, 64, 120, 160]               0\n",
      "         Identity-71         [-1, 64, 120, 160]               0\n",
      "        Attention-72         [-1, 64, 120, 160]               0\n",
      "     DecoderBlock-73         [-1, 64, 120, 160]               0\n",
      "         Identity-74        [-1, 128, 240, 320]               0\n",
      "        Attention-75        [-1, 128, 240, 320]               0\n",
      "           Conv2d-76         [-1, 32, 240, 320]          36,864\n",
      "      BatchNorm2d-77         [-1, 32, 240, 320]              64\n",
      "             ReLU-78         [-1, 32, 240, 320]               0\n",
      "           Conv2d-79         [-1, 32, 240, 320]           9,216\n",
      "      BatchNorm2d-80         [-1, 32, 240, 320]              64\n",
      "             ReLU-81         [-1, 32, 240, 320]               0\n",
      "         Identity-82         [-1, 32, 240, 320]               0\n",
      "        Attention-83         [-1, 32, 240, 320]               0\n",
      "     DecoderBlock-84         [-1, 32, 240, 320]               0\n",
      "           Conv2d-85         [-1, 16, 480, 640]           4,608\n",
      "      BatchNorm2d-86         [-1, 16, 480, 640]              32\n",
      "             ReLU-87         [-1, 16, 480, 640]               0\n",
      "           Conv2d-88         [-1, 16, 480, 640]           2,304\n",
      "      BatchNorm2d-89         [-1, 16, 480, 640]              32\n",
      "             ReLU-90         [-1, 16, 480, 640]               0\n",
      "         Identity-91         [-1, 16, 480, 640]               0\n",
      "        Attention-92         [-1, 16, 480, 640]               0\n",
      "     DecoderBlock-93         [-1, 16, 480, 640]               0\n",
      "      UnetDecoder-94         [-1, 16, 480, 640]               0\n",
      "           Conv2d-95         [-1, 11, 480, 640]           1,595\n",
      "         Identity-96         [-1, 11, 480, 640]               0\n",
      "         Identity-97         [-1, 11, 480, 640]               0\n",
      "       Activation-98         [-1, 11, 480, 640]               0\n",
      "================================================================\n",
      "Total params: 8,058,939\n",
      "Trainable params: 8,058,939\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 3.52\n",
      "Forward/backward pass size (MB): 1305.47\n",
      "Params size (MB): 30.74\n",
      "Estimated Total Size (MB): 1339.73\n",
      "----------------------------------------------------------------\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "unet = models.unet.Unet('resnet10', classes = 11, encoder_weights = None).cuda()\n",
    "# print(unet)\n",
    "print(summary(unet, input_size = (3, 480, 640)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SaveFeatures :\n",
    "    def __init__(self, m) : \n",
    "        self.handle = m.register_forward_hook(self.hook_fn)\n",
    "    def hook_fn(self, m, inp, outp) : \n",
    "        self.features = outp\n",
    "    def remove(self) :\n",
    "        self.handle.remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "sf = [SaveFeatures(m) for m in [unet.encoder.relu, \n",
    "                                unet.encoder.layer1, \n",
    "                                unet.encoder.layer2, \n",
    "                                unet.encoder.layer3, \n",
    "                                unet.encoder.layer4, \n",
    "                                unet.decoder.blocks[0], \n",
    "                                unet.decoder.blocks[1], \n",
    "                                unet.decoder.blocks[2], \n",
    "                                unet.decoder.blocks[3], \n",
    "                                unet.decoder.blocks[4] \n",
    "                               ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean = [0.35675976, 0.37380189, 0.3764753], std = [0.32064945, 0.32098866, 0.32325324])\n",
    "])\n",
    "trainset = KITTI(split = 'train', transform = transform)\n",
    "testset = KITTI(split = 'test', transform = transform)\n",
    "\n",
    "trainloader = DataLoader(trainset, batch_size = 2, shuffle = True)\n",
    "testloader = DataLoader(testset, batch_size = 1, shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 64, 128, 512])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 512, 8, 32])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 32, 128, 512])\n",
      "torch.Size([2, 16, 256, 1024])\n"
     ]
    }
   ],
   "source": [
    "# x, _ = next(iter(trainloader))\n",
    "x = torch.randn(2, 3, 256, 1024)\n",
    "x = x.cuda()\n",
    "out = unet(x)\n",
    "for i in range(10) :\n",
    "    print(sf[i].features.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 256, 256]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 256, 256]             128\n",
      "              ReLU-3         [-1, 64, 256, 256]               0\n",
      "         MaxPool2d-4         [-1, 64, 128, 128]               0\n",
      "            Conv2d-5         [-1, 64, 128, 128]          36,864\n",
      "       BatchNorm2d-6         [-1, 64, 128, 128]             128\n",
      "              ReLU-7         [-1, 64, 128, 128]               0\n",
      "            Conv2d-8         [-1, 64, 128, 128]          36,864\n",
      "       BatchNorm2d-9         [-1, 64, 128, 128]             128\n",
      "             ReLU-10         [-1, 64, 128, 128]               0\n",
      "       BasicBlock-11         [-1, 64, 128, 128]               0\n",
      "           Conv2d-12         [-1, 64, 128, 128]          36,864\n",
      "      BatchNorm2d-13         [-1, 64, 128, 128]             128\n",
      "             ReLU-14         [-1, 64, 128, 128]               0\n",
      "           Conv2d-15         [-1, 64, 128, 128]          36,864\n",
      "      BatchNorm2d-16         [-1, 64, 128, 128]             128\n",
      "             ReLU-17         [-1, 64, 128, 128]               0\n",
      "       BasicBlock-18         [-1, 64, 128, 128]               0\n",
      "           Conv2d-19          [-1, 128, 64, 64]          73,728\n",
      "      BatchNorm2d-20          [-1, 128, 64, 64]             256\n",
      "             ReLU-21          [-1, 128, 64, 64]               0\n",
      "           Conv2d-22          [-1, 128, 64, 64]         147,456\n",
      "      BatchNorm2d-23          [-1, 128, 64, 64]             256\n",
      "           Conv2d-24          [-1, 128, 64, 64]           8,192\n",
      "      BatchNorm2d-25          [-1, 128, 64, 64]             256\n",
      "             ReLU-26          [-1, 128, 64, 64]               0\n",
      "       BasicBlock-27          [-1, 128, 64, 64]               0\n",
      "           Conv2d-28          [-1, 128, 64, 64]         147,456\n",
      "      BatchNorm2d-29          [-1, 128, 64, 64]             256\n",
      "             ReLU-30          [-1, 128, 64, 64]               0\n",
      "           Conv2d-31          [-1, 128, 64, 64]         147,456\n",
      "      BatchNorm2d-32          [-1, 128, 64, 64]             256\n",
      "             ReLU-33          [-1, 128, 64, 64]               0\n",
      "       BasicBlock-34          [-1, 128, 64, 64]               0\n",
      "           Conv2d-35          [-1, 256, 32, 32]         294,912\n",
      "      BatchNorm2d-36          [-1, 256, 32, 32]             512\n",
      "             ReLU-37          [-1, 256, 32, 32]               0\n",
      "           Conv2d-38          [-1, 256, 32, 32]         589,824\n",
      "      BatchNorm2d-39          [-1, 256, 32, 32]             512\n",
      "           Conv2d-40          [-1, 256, 32, 32]          32,768\n",
      "      BatchNorm2d-41          [-1, 256, 32, 32]             512\n",
      "             ReLU-42          [-1, 256, 32, 32]               0\n",
      "       BasicBlock-43          [-1, 256, 32, 32]               0\n",
      "           Conv2d-44          [-1, 256, 32, 32]         589,824\n",
      "      BatchNorm2d-45          [-1, 256, 32, 32]             512\n",
      "             ReLU-46          [-1, 256, 32, 32]               0\n",
      "           Conv2d-47          [-1, 256, 32, 32]         589,824\n",
      "      BatchNorm2d-48          [-1, 256, 32, 32]             512\n",
      "             ReLU-49          [-1, 256, 32, 32]               0\n",
      "       BasicBlock-50          [-1, 256, 32, 32]               0\n",
      "           Conv2d-51          [-1, 512, 16, 16]       1,179,648\n",
      "      BatchNorm2d-52          [-1, 512, 16, 16]           1,024\n",
      "             ReLU-53          [-1, 512, 16, 16]               0\n",
      "           Conv2d-54          [-1, 512, 16, 16]       2,359,296\n",
      "      BatchNorm2d-55          [-1, 512, 16, 16]           1,024\n",
      "           Conv2d-56          [-1, 512, 16, 16]         131,072\n",
      "      BatchNorm2d-57          [-1, 512, 16, 16]           1,024\n",
      "             ReLU-58          [-1, 512, 16, 16]               0\n",
      "       BasicBlock-59          [-1, 512, 16, 16]               0\n",
      "           Conv2d-60          [-1, 512, 16, 16]       2,359,296\n",
      "      BatchNorm2d-61          [-1, 512, 16, 16]           1,024\n",
      "             ReLU-62          [-1, 512, 16, 16]               0\n",
      "           Conv2d-63          [-1, 512, 16, 16]       2,359,296\n",
      "      BatchNorm2d-64          [-1, 512, 16, 16]           1,024\n",
      "             ReLU-65          [-1, 512, 16, 16]               0\n",
      "       BasicBlock-66          [-1, 512, 16, 16]               0\n",
      "    ResNetEncoder-67  [[-1, 3, 512, 512], [-1, 64, 256, 256], [-1, 64, 128, 128], [-1, 128, 64, 64], [-1, 256, 32, 32], [-1, 512, 16, 16]]               0\n",
      "         Identity-68          [-1, 512, 16, 16]               0\n",
      "         Identity-69          [-1, 768, 32, 32]               0\n",
      "        Attention-70          [-1, 768, 32, 32]               0\n",
      "           Conv2d-71          [-1, 256, 32, 32]       1,769,472\n",
      "      BatchNorm2d-72          [-1, 256, 32, 32]             512\n",
      "             ReLU-73          [-1, 256, 32, 32]               0\n",
      "           Conv2d-74          [-1, 256, 32, 32]         589,824\n",
      "      BatchNorm2d-75          [-1, 256, 32, 32]             512\n",
      "             ReLU-76          [-1, 256, 32, 32]               0\n",
      "         Identity-77          [-1, 256, 32, 32]               0\n",
      "        Attention-78          [-1, 256, 32, 32]               0\n",
      "     DecoderBlock-79          [-1, 256, 32, 32]               0\n",
      "         Identity-80          [-1, 384, 64, 64]               0\n",
      "        Attention-81          [-1, 384, 64, 64]               0\n",
      "           Conv2d-82          [-1, 128, 64, 64]         442,368\n",
      "      BatchNorm2d-83          [-1, 128, 64, 64]             256\n",
      "             ReLU-84          [-1, 128, 64, 64]               0\n",
      "           Conv2d-85          [-1, 128, 64, 64]         147,456\n",
      "      BatchNorm2d-86          [-1, 128, 64, 64]             256\n",
      "             ReLU-87          [-1, 128, 64, 64]               0\n",
      "         Identity-88          [-1, 128, 64, 64]               0\n",
      "        Attention-89          [-1, 128, 64, 64]               0\n",
      "     DecoderBlock-90          [-1, 128, 64, 64]               0\n",
      "         Identity-91        [-1, 192, 128, 128]               0\n",
      "        Attention-92        [-1, 192, 128, 128]               0\n",
      "           Conv2d-93         [-1, 64, 128, 128]         110,592\n",
      "      BatchNorm2d-94         [-1, 64, 128, 128]             128\n",
      "             ReLU-95         [-1, 64, 128, 128]               0\n",
      "           Conv2d-96         [-1, 64, 128, 128]          36,864\n",
      "      BatchNorm2d-97         [-1, 64, 128, 128]             128\n",
      "             ReLU-98         [-1, 64, 128, 128]               0\n",
      "         Identity-99         [-1, 64, 128, 128]               0\n",
      "       Attention-100         [-1, 64, 128, 128]               0\n",
      "    DecoderBlock-101         [-1, 64, 128, 128]               0\n",
      "        Identity-102        [-1, 128, 256, 256]               0\n",
      "       Attention-103        [-1, 128, 256, 256]               0\n",
      "          Conv2d-104         [-1, 32, 256, 256]          36,864\n",
      "     BatchNorm2d-105         [-1, 32, 256, 256]              64\n",
      "            ReLU-106         [-1, 32, 256, 256]               0\n",
      "          Conv2d-107         [-1, 32, 256, 256]           9,216\n",
      "     BatchNorm2d-108         [-1, 32, 256, 256]              64\n",
      "            ReLU-109         [-1, 32, 256, 256]               0\n",
      "        Identity-110         [-1, 32, 256, 256]               0\n",
      "       Attention-111         [-1, 32, 256, 256]               0\n",
      "    DecoderBlock-112         [-1, 32, 256, 256]               0\n",
      "          Conv2d-113         [-1, 16, 512, 512]           4,608\n",
      "     BatchNorm2d-114         [-1, 16, 512, 512]              32\n",
      "            ReLU-115         [-1, 16, 512, 512]               0\n",
      "          Conv2d-116         [-1, 16, 512, 512]           2,304\n",
      "     BatchNorm2d-117         [-1, 16, 512, 512]              32\n",
      "            ReLU-118         [-1, 16, 512, 512]               0\n",
      "        Identity-119         [-1, 16, 512, 512]               0\n",
      "       Attention-120         [-1, 16, 512, 512]               0\n",
      "    DecoderBlock-121         [-1, 16, 512, 512]               0\n",
      "     UnetDecoder-122         [-1, 16, 512, 512]               0\n",
      "          Conv2d-123         [-1, 19, 512, 512]           2,755\n",
      "        Identity-124         [-1, 19, 512, 512]               0\n",
      "         Softmax-125         [-1, 19, 512, 512]               0\n",
      "      Activation-126         [-1, 19, 512, 512]               0\n",
      "================================================================\n",
      "Total params: 14,330,819\n",
      "Trainable params: 14,330,819\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 3.00\n",
      "Forward/backward pass size (MB): 1283.00\n",
      "Params size (MB): 54.67\n",
      "Estimated Total Size (MB): 1340.67\n",
      "----------------------------------------------------------------\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "unet = models.unet.Unet('resnet18', classes = 19, activation = 'softmax', encoder_weights = None).cuda()\n",
    "# print(unet)\n",
    "print(summary(unet, input_size = (3, 512, 512)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 64, 128, 512])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 512, 8, 32])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 32, 128, 512])\n",
      "torch.Size([2, 16, 256, 1024])\n"
     ]
    }
   ],
   "source": [
    "# x, _ = next(iter(trainloader))\n",
    "x = torch.randn(2, 3, 256, 1024)\n",
    "x = x.cuda()\n",
    "out = unet(x)\n",
    "for i in range(10) :\n",
    "    print(sf[i].features.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "unet = models.unet.Unet('resnet34', classes = 19, activation = 'softmax').cuda()\n",
    "unet2 = models.unet.Unet('resnet18', classes = 19, activation = 'softmax').cuda()\n",
    "\n",
    "sf = [SaveFeatures(m) for m in [unet.encoder.relu, \n",
    "                                unet.encoder.layer1, \n",
    "                                unet.encoder.layer2, \n",
    "                                unet.encoder.layer3, \n",
    "                                unet.encoder.layer4, \n",
    "                                unet.decoder.blocks[0], \n",
    "                                unet.decoder.blocks[1], \n",
    "                                unet.decoder.blocks[2], \n",
    "                                unet.decoder.blocks[3], \n",
    "                                unet.decoder.blocks[4] \n",
    "                               ]]\n",
    "\n",
    "sf2 = [SaveFeatures(m) for m in [unet2.encoder.relu, \n",
    "                                unet2.encoder.layer1, \n",
    "                                unet2.encoder.layer2, \n",
    "                                unet2.encoder.layer3, \n",
    "                                unet2.encoder.layer4, \n",
    "                                unet2.decoder.blocks[0], \n",
    "                                unet2.decoder.blocks[1], \n",
    "                                unet2.decoder.blocks[2], \n",
    "                                unet2.decoder.blocks[3], \n",
    "                                unet2.decoder.blocks[4] \n",
    "                               ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 64, 128, 512])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 512, 8, 32])\n",
      "torch.Size([2, 256, 16, 64])\n",
      "torch.Size([2, 128, 32, 128])\n",
      "torch.Size([2, 64, 64, 256])\n",
      "torch.Size([2, 32, 128, 512])\n",
      "torch.Size([2, 16, 256, 1024])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn(2, 3, 256, 1024).cuda()\n",
    "_ = unet(x)\n",
    "_ = unet2(x)\n",
    "for i in range(10) :\n",
    "    print(sf[i].features.shape)\n",
    "    assert(sf[i].features.shape == sf2[i].features.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (pyt)",
   "language": "python",
   "name": "pyt"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
